# 使用AI辅助
# Gitee Repo

import torch
import numpy as np
from vpython import *

torch.set_default_dtype(torch.float64)

# 设置参数
L = 5
N = 10  # 粒子个数
dt = 0.01

allow_friction = True
allow_boxed = True

# 初始化粒子位置和速度 (使用PyTorch张量)
r0 = (2 * torch.rand(N, 3) - 1) * L
v0 = 2 * torch.rand(N, 3) - 1

# 居中处理
r0 -= torch.mean(r0, dim=0)
v0 -= torch.mean(v0, dim=0)

# LJ-like 势
def compute_f(r):
    if r.ndim != 2 or r.shape[0] != N or r.shape[1] != 3:
        raise ValueError('123')
        
    rvec = r[None,:,:]-r[:,None,:] #(N,N,3)
    dist = torch.norm(rvec,dim=2) #(N,N)
    fmag = 1/dist**3 - 1/dist**6
    fmag[dist<0.0001] = 0
    fmat = fmag[:,:,None]*rvec #(N,N,3)
    f = torch.sum(fmat,dim=1) #(N,3)
    return f

# 创建VPython场景和粒子
scene = canvas(title='N-body Simulation', width=800, height=600, range=L)
scene.background = color.white 
particles = []
for i in range(N):
    p = sphere(pos=vector(r0[i,0].item(),r0[i,1].item(),r0[i,2].item()), radius=0.2, color=color.orange)
    particles.append(p)

# 添加边界显示
box = box(pos=vector(0,0,0), size=vector(2*L,2*L,2*L), opacity=0.2)

# 主循环
for tick in range(10000):
    rate(60)  # 控制动画速度
    
    # 计算半步长速度
    f = compute_f(r0)
    vhalf = v0 + 0.5 * f * dt
    
    # 计算下一时刻位置
    r1 = r0 + vhalf * dt
    
    # 计算下一时刻速度
    f = compute_f(r1)
    v1 = vhalf + 0.5 * f * dt
    
    # 更新位置与速度
    r0, v0 = r1, v1
    
    # 摩擦
    if allow_friction:
        v0 *= 0.998
    
    # 刚性边界
    if allow_boxed:
        # 检测是否越界并反弹
        below_lower = r0 < -L
        above_upper = r0 > L
        
        # 应用边界条件
        r0[below_lower] = -L
        r0[above_upper] = L
        
        # 反转速度
        v0[below_lower] = -v0[below_lower]
        v0[above_upper] = -v0[above_upper]
        
    
    # 检查是否发散
    if torch.any(torch.isnan(r0)) or torch.any(torch.isinf(r0)):
        print("Simulation diverged!")
        break
    
    # 检查是否所有粒子都停止了
    if torch.all(torch.norm(v0,dim=1) < 0.01):
        print("All particles have stopped")
        break
    
    # 更新VPython显示
    if tick % 5 == 0:
        for i in range(N):
            particles[i].pos = vector(r0[i,0].item(),r0[i,1].item(),r0[i,2].item())
            
        Ptot = torch.norm(torch.sum(v0,dim=0)) # 总动量，若无刚性边界则应该守恒
        Ltot = torch.norm(torch.sum(torch.linalg.cross(r0,v0),dim=0)) # 总角动量，若无刚性边界则应该守恒
        print(f'Tick:{tick},Ptot:{Ptot:.2f},Ltot:{Ltot:.2f}')

print('Sim finished')
print('Press Enter to Exit')
input()
exit()
